package com.zbd.microservices.service;

import java.io.UnsupportedEncodingException;
import java.math.BigInteger;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

import javax.inject.Inject;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.domain.Example;

import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.common.util.SerializationUtils;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import com.zbd.microservices.domain.OauthAccessToken;
import com.zbd.microservices.domain.OauthRefreshToken;
import com.zbd.microservices.repository.OauthAccessTokenRepository;
import com.zbd.microservices.repository.OauthRefreshTokenRepository;

@Service
@Transactional
public class HibernateTokenStoreService implements TokenStore
{
	private static final Logger LOG = LoggerFactory.getLogger(HibernateTokenStoreService.class);

	private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();

	@Inject
	private OauthAccessTokenRepository oauthAccessTokenRepository;

	@Inject
	private OauthRefreshTokenRepository oauthRefreshTokenRepository;

	@Override
	public OAuth2Authentication readAuthentication(OAuth2AccessToken token)
	{
		return readAuthentication(token.getValue());
	}

	@Override
	public OAuth2Authentication readAuthentication(String token)
	{
		OAuth2Authentication authentication = null;

		try
		{
			String tokenKey = extractTokenKey(token);

			OauthAccessToken oauthAccessToken4Query = new OauthAccessToken();
			oauthAccessToken4Query.setTokenId(tokenKey);
			Example<OauthAccessToken> example = Example.of(oauthAccessToken4Query);

			OauthAccessToken oauthAccessToken = oauthAccessTokenRepository.findOne(example);
			if(oauthAccessToken != null)
			{
				authentication = deserializeAuthentication(oauthAccessToken.getAuthentication());
			}
			else
			{
				LOG.info("Failed to find access token for token {}", token);
			}
		}
		catch (IllegalArgumentException e)
		{
			LOG.warn("Failed to deserialize authentication for " + token, e);
			removeAccessToken(token);
		}

		return authentication;
	}

	@Override
	public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication)
	{
		String refreshToken = null;
		if (token.getRefreshToken() != null)
		{
			refreshToken = token.getRefreshToken().getValue();
		}

		if (readAccessToken(token.getValue()) != null)
		{
			removeAccessToken(token.getValue());
		}

		OauthAccessToken oauthAccessToken = new OauthAccessToken();

		oauthAccessToken.setTokenId(extractTokenKey(token.getValue()));
		oauthAccessToken.setToken(serializeAccessToken(token));
		oauthAccessToken.setAuthenticationId(authenticationKeyGenerator.extractKey(authentication));
		oauthAccessToken.setUserName(authentication.isClientOnly() ? null : authentication.getName());
		oauthAccessToken.setClientId(authentication.getOAuth2Request().getClientId());
		oauthAccessToken.setAuthentication(serializeAuthentication(authentication));
		oauthAccessToken.setRefreshToken(extractTokenKey(refreshToken));

		oauthAccessTokenRepository.save(oauthAccessToken);
	}

	@Override
	public OAuth2AccessToken readAccessToken(String tokenValue)
	{
		OAuth2AccessToken accessToken = null;

		try
		{
			String tokenId = extractTokenKey(tokenValue);

			OauthAccessToken oauthAccessToken4Query = new OauthAccessToken();
			oauthAccessToken4Query.setTokenId(tokenId);
			Example<OauthAccessToken> example = Example.of(oauthAccessToken4Query);

			OauthAccessToken oauthAccessToken = oauthAccessTokenRepository.findOne(example);
			if(oauthAccessToken != null)
			{
				accessToken = deserializeAccessToken(oauthAccessToken.getToken());
			}
			else
			{
				LOG.info("Failed to find access token for token {}", tokenValue);
			}
		}
		catch (IllegalArgumentException e)
		{
			LOG.warn("Failed to deserialize access token for " + tokenValue, e);
			removeAccessToken(tokenValue);
		}

		return accessToken;
	}

	public void removeAccessToken(String tokenValue)
	{
		oauthAccessTokenRepository.delete(extractTokenKey(tokenValue));
	}

	@Override
	public void removeAccessToken(OAuth2AccessToken accessToken)
	{
		removeAccessToken(accessToken.getValue());
	}

	@Override
	public void storeRefreshToken(OAuth2RefreshToken refreshToken, OAuth2Authentication authentication)
	{
		OauthRefreshToken oauthRefreshToken = new OauthRefreshToken();

		oauthRefreshToken.setTokenId(refreshToken.getValue());
		oauthRefreshToken.setToken(serializeRefreshToken(refreshToken));
		oauthRefreshToken.setAuthentication(serializeAuthentication(authentication));

		oauthRefreshTokenRepository.save(oauthRefreshToken);
	}

	@Override
	public OAuth2RefreshToken readRefreshToken(String token)
	{
		OAuth2RefreshToken refreshToken = null;

		try
		{
			String tokenKey = extractTokenKey(token);

			OauthRefreshToken oauthRefreshToken4Query = new OauthRefreshToken();
			oauthRefreshToken4Query.setTokenId(tokenKey);
			Example<OauthRefreshToken> example = Example.of(oauthRefreshToken4Query);

			OauthRefreshToken oauthRefreshToken = oauthRefreshTokenRepository.findOne(example);
			
			if (oauthRefreshToken != null)
			{
				refreshToken = deserializeRefreshToken(oauthRefreshToken.getToken());
			}
			else
			{
				LOG.info("Failed to find refresh token for token {}", token);
			}
		}
		catch (IllegalArgumentException e)
		{
			LOG.warn("Failed to deserialize refresh token for token " + token, e);
			removeRefreshToken(token);
		}

		return refreshToken;
	}

	@Override
	public OAuth2Authentication readAuthenticationForRefreshToken(OAuth2RefreshToken token)
	{
		return readAuthenticationForRefreshToken(token.getValue());
	}

	public OAuth2Authentication readAuthenticationForRefreshToken(String token)
	{
		OAuth2Authentication authentication = null;

		try
		{
			String tokenKey = extractTokenKey(token);

			OauthRefreshToken oauthRefreshToken4Query = new OauthRefreshToken();
			oauthRefreshToken4Query.setTokenId(tokenKey);
			Example<OauthRefreshToken> example = Example.of(oauthRefreshToken4Query);

			OauthRefreshToken oauthRefreshToken = oauthRefreshTokenRepository.findOne(example);
			if(oauthRefreshToken != null)
			{
				authentication = deserializeAuthentication(oauthRefreshToken.getToken());
			}
			else
			{
				LOG.info("Failed to find access token for token {}", token);
			}
		}
		catch (IllegalArgumentException e)
		{
			LOG.warn("Failed to deserialize access token for " + token, e);
			removeRefreshToken(token);
		}

		return authentication;
	}

	@Override
	public void removeRefreshToken(OAuth2RefreshToken token)
	{
		removeRefreshToken(token.getValue());
	}

	public void removeRefreshToken(String token)
	{
		String tokenId = extractTokenKey(token);
		OauthRefreshToken oauthRefreshToken4Query = new OauthRefreshToken();
		oauthRefreshToken4Query.setTokenId(tokenId);
		Example<OauthRefreshToken> example = Example.of(oauthRefreshToken4Query);
		
		OauthRefreshToken oauthRefreshToken4Delete = oauthRefreshTokenRepository.findOne(example);
		
		if(oauthRefreshToken4Delete != null)
		{
			oauthRefreshTokenRepository.delete(oauthRefreshToken4Delete);
		}
		else
		{
			LOG.info("Failed to find access token for token {}", token);
		}
	}

	@Override
	public void removeAccessTokenUsingRefreshToken(OAuth2RefreshToken refreshToken)
	{
		removeAccessTokenUsingRefreshToken(refreshToken.getValue());
	}

	public void removeAccessTokenUsingRefreshToken(String refreshToken)
	{
		OauthAccessToken oauthAccessToken4Query = new OauthAccessToken();
		oauthAccessToken4Query.setRefreshToken(refreshToken);
		Example<OauthAccessToken> example = Example.of(oauthAccessToken4Query);

		OauthAccessToken oauthAccessToken4Delete = oauthAccessTokenRepository.findOne(example);
		if(oauthAccessToken4Delete != null)
		{
			oauthAccessTokenRepository.delete(oauthAccessToken4Delete);
		}
		else
		{
			LOG.info("Failed to find access token for refresh token {}", refreshToken);
		}
	}

	@Override
	public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication)
	{
		OAuth2AccessToken accessToken = null;

		String key = authenticationKeyGenerator.extractKey(authentication);
		try
		{
			OauthAccessToken oauthAccessToken = oauthAccessTokenRepository.findOne(key);
			if(oauthAccessToken != null)
			{
				accessToken = deserializeAccessToken(oauthAccessToken.getToken());
			}
			else
			{
				LOG.debug("Failed to find access token for authentication " + authentication);
			}
		}
		catch (IllegalArgumentException e)
		{
			LOG.error("Could not extract access token for authentication " + authentication, e);
		}

		if (accessToken != null && !key.equals(authenticationKeyGenerator.extractKey(readAuthentication(accessToken.getValue()))))
		{
			removeAccessToken(accessToken.getValue());
			storeAccessToken(accessToken, authentication);
		}
		return accessToken;
	}

	@Override
	public Collection<OAuth2AccessToken> findTokensByClientIdAndUserName(String clientId, String userName)
	{
		List<OAuth2AccessToken> accessTokens = new ArrayList<OAuth2AccessToken>();
		OauthAccessToken oauthAccessToken4Query = new OauthAccessToken();
		oauthAccessToken4Query.setClientId(clientId);
		oauthAccessToken4Query.setUserName(userName);
		Example<OauthAccessToken> example = Example.of(oauthAccessToken4Query);

		List<OauthAccessToken> oauthAccessTokens = oauthAccessTokenRepository.findAll(example);
		if (oauthAccessTokens != null && oauthAccessTokens.size() > 0)
		{
			for (OauthAccessToken oauthAccessToken : oauthAccessTokens)
			{
				try
				{
					accessTokens.add(deserializeAccessToken(oauthAccessToken.getToken()));
				}
				catch (IllegalArgumentException e)
				{

					LOG.warn("Delete illegal token authentication_id={}, client_id={}, user_name={}", oauthAccessToken.getAuthenticationId(), oauthAccessToken.getClientId(),
							oauthAccessToken.getUserName());
					oauthAccessTokenRepository.delete(oauthAccessToken);
				}
			}
		}
		else
		{
			LOG.info("Failed to find access token for client_id = {} and user_name = {} ", clientId, userName);
		}
		accessTokens = removeNulls(accessTokens);

		return accessTokens;
	}

	@Override
	public Collection<OAuth2AccessToken> findTokensByClientId(String clientId)
	{
		List<OAuth2AccessToken> accessTokens = new ArrayList<OAuth2AccessToken>();

		OauthAccessToken oauthAccessToken4Query = new OauthAccessToken();
		oauthAccessToken4Query.setClientId(clientId);
		Example<OauthAccessToken> example = Example.of(oauthAccessToken4Query);

		List<OauthAccessToken> oauthAccessTokens = oauthAccessTokenRepository.findAll(example);
		if (oauthAccessTokens != null && oauthAccessTokens.size() > 0)
		{
			for (OauthAccessToken oauthAccessToken : oauthAccessTokens)
			{
				try
				{
					accessTokens.add(deserializeAccessToken(oauthAccessToken.getToken()));
				}
				catch (IllegalArgumentException e)
				{
					LOG.warn("Delete illegal token authentication_id={}, client_id={}", oauthAccessToken.getAuthenticationId(), oauthAccessToken.getClientId());
					oauthAccessTokenRepository.delete(oauthAccessToken);
				}
			}
		}
		else
		{
			LOG.info("Failed to find access token for clientId {}", clientId);
		}
		accessTokens = removeNulls(accessTokens);

		return accessTokens;
	}

	protected byte[] serializeAccessToken(OAuth2AccessToken token)
	{
		return SerializationUtils.serialize(token);
	}

	protected byte[] serializeRefreshToken(OAuth2RefreshToken token)
	{
		return SerializationUtils.serialize(token);
	}

	protected byte[] serializeAuthentication(OAuth2Authentication authentication)
	{
		return SerializationUtils.serialize(authentication);
	}

	protected OAuth2AccessToken deserializeAccessToken(byte[] token)
	{
		return SerializationUtils.deserialize(token);
	}

	protected OAuth2RefreshToken deserializeRefreshToken(byte[] token)
	{
		return SerializationUtils.deserialize(token);
	}

	protected OAuth2Authentication deserializeAuthentication(byte[] authentication)
	{
		return SerializationUtils.deserialize(authentication);
	}

	protected String extractTokenKey(String value)
	{
		if (value == null)
		{
			return null;
		}
		MessageDigest digest;
		try
		{
			digest = MessageDigest.getInstance("MD5");
		}
		catch (NoSuchAlgorithmException e)
		{
			throw new IllegalStateException("MD5 algorithm not available.  Fatal (should be in the JDK).");
		}

		try
		{
			byte[] bytes = digest.digest(value.getBytes("UTF-8"));
			return String.format("%032x", new BigInteger(1, bytes));
		}
		catch (UnsupportedEncodingException e)
		{
			throw new IllegalStateException("UTF-8 encoding not available.  Fatal (should be in the JDK).");
		}
	}

	private List<OAuth2AccessToken> removeNulls(List<OAuth2AccessToken> accessTokens)
	{
		List<OAuth2AccessToken> tokens = new ArrayList<OAuth2AccessToken>();
		for (OAuth2AccessToken token : accessTokens)
		{
			if (token != null)
			{
				tokens.add(token);
			}
		}
		return tokens;
	}
}
